#include <pcl/point_types.h>
#include <pcl/io/pcd_io.h>
#include <pcl/visualization/pcl_visualizer.h>
#include <vector>
#include <string>
#include <iostream>
#include <thread>
#include <chrono>
#include <atomic>
#include <filesystem>
#include <X11/Xlib.h>
#include <numeric>
#include <tuple>
#include <cmath>

// 假设自定义头文件
#include <ground_remove.h>
#include <component_clustering.h>
#include <box_fitting.h>

// ------ 常量和类型 ------
namespace fs = std::filesystem;
const std::string INPUT_FOLDER  = "/home/chenhao/kitti3d_segmentation_benchmark/data/pcd";
const std::string OUTPUT_FOLDER = "/home/chenhao/kitti3d_segmentation_benchmark/data/visualization_output";
const uint8_t GROUND_R = 50, GROUND_G = 205, GROUND_B = 50;
const uint8_t BOX_R = 255, BOX_G = 0, BOX_B = 0;

// ------ 全局状态 ------
std::atomic<size_t> current_idx(0);
std::atomic<bool> next_flag(false);
std::atomic<bool> prev_flag(false);
std::atomic<bool> quit(false);
std::atomic<bool> refresh(false);

// ------ 键盘响应 ------
void keyboardEventOccurred(const pcl::visualization::KeyboardEvent& event, void* viewer_ptr)
{
    if (event.keyDown()) {
        if (event.getKeySym() == "n" || event.getKeySym() == "Right") {
            next_flag.store(true);
        }
        else if (event.getKeySym() == "p" || event.getKeySym() == "Left") {
            prev_flag.store(true);
        }
        else if (event.getKeySym() == "q" || event.getKeySym() == "Escape") {
            quit.store(true);
        }
        else if (event.getKeySym() == "r") {
            refresh.store(true);
        }
    }
}

// ------ 聚类着色 ------
std::tuple<uint8_t, uint8_t, uint8_t> generateClusterColor(size_t idx, size_t total) {
    float ratio = total > 1 ? static_cast<float>(idx) / (total-1) : 0;
    return {
        static_cast<uint8_t>(127.5f * (1 + std::sin(2.f * M_PI * ratio))),
        static_cast<uint8_t>(127.5f * (1 + std::sin(2.f * M_PI * ratio + 2.f))),
        static_cast<uint8_t>(127.5f * (1 + std::sin(2.f * M_PI * ratio + 4.f)))
    };
}

// ------ 可视化包围盒 ------
void visualizeBoundingBox(
    const pcl::visualization::PCLVisualizer::Ptr& viewer,
    const pcl::PointCloud<pcl::PointXYZI>& box_points,
    const std::tuple<uint8_t, uint8_t, uint8_t>& color,
    const std::string& box_name)
{
    if (box_points.size() != 8) return;
    const auto& pt = box_points.points;
    const auto [r, g, b] = color;
    auto addLine = [&](int p1, int p2, int idx) {
        viewer->addLine(pt[p1], pt[p2], r/255.0f, g/255.0f, b/255.0f, box_name + "_line" + std::to_string(idx));
    };
    // 下底面
    addLine(0, 1, 1); addLine(1, 2, 2); addLine(2, 3, 3); addLine(3, 0, 4);
    // 上底面
    addLine(4, 5, 5); addLine(5, 6, 6); addLine(6, 7, 7); addLine(7, 4, 8);
    // 侧面
    addLine(0, 4, 9); addLine(1, 5, 10); addLine(2, 6, 11); addLine(3, 7, 12);
}

// ------ 保存地面点 ------
void saveGround(
    const pcl::PointCloud<pcl::PointXYZI>::Ptr& ground,
    size_t frame_idx)
{
    try {
        if (!fs::exists(OUTPUT_FOLDER)) fs::create_directories(OUTPUT_FOLDER);
        pcl::PointCloud<pcl::PointXYZRGB>::Ptr out_cloud(new pcl::PointCloud<pcl::PointXYZRGB>);
        out_cloud->reserve(ground->size());
        for (const auto& pt : *ground) {
            pcl::PointXYZRGB p; p.x = pt.x; p.y = pt.y; p.z = pt.z;
            p.r = GROUND_R; p.g = GROUND_G; p.b = GROUND_B;
            out_cloud->push_back(p);
        }
        std::string out_path = OUTPUT_FOLDER + "/frame_ground_" + std::to_string(frame_idx) + ".pcd";
        if (pcl::io::savePCDFileBinary(out_path, *out_cloud) == -1) throw std::runtime_error("Failed to save: " + out_path);
        std::cout << "已保存地面点云: " << out_path << std::endl;
    } catch(const std::exception& e) { std::cerr << "地面点保存出错: " << e.what() << std::endl; }
}

// ------ 保存地物点 ------
void saveObjects(
    const std::vector<pcl::PointCloud<pcl::PointXYZI>>& objects,
    size_t frame_idx)
{
    try {
        if (!fs::exists(OUTPUT_FOLDER)) fs::create_directories(OUTPUT_FOLDER);
        pcl::PointCloud<pcl::PointXYZRGB>::Ptr out_cloud(new pcl::PointCloud<pcl::PointXYZRGB>);
        size_t n = 0; for (auto& x : objects) n += x.size();
        out_cloud->reserve(n);
        for (size_t i = 0; i < objects.size(); ++i) {
            if (objects[i].empty()) continue;
            auto [r, g, b] = generateClusterColor(i, objects.size());
            for (const auto& pt : objects[i]) {
                pcl::PointXYZRGB p; p.x = pt.x; p.y = pt.y; p.z = pt.z;
                p.r = r; p.g = g; p.b = b;
                out_cloud->push_back(p);
            }
        }
        std::string out_path = OUTPUT_FOLDER + "/frame_objects_" + std::to_string(frame_idx) + ".pcd";
        if (pcl::io::savePCDFileBinary(out_path, *out_cloud) == -1) throw std::runtime_error("Failed to save: " + out_path);
        std::cout << "已保存地物点云: " << out_path << std::endl;
    } catch(const std::exception& e) { std::cerr << "地物点保存出错: " << e.what() << std::endl; }
}

// ------ 帮助信息 ------
void displayControlInfo() {
    if (refresh.load() || current_idx == 0) {
        std::cout << "\n===== 控制指令 =====" << std::endl
                  << "n/→  : 下一帧" << std::endl
                  << "p/←  : 上一帧" << std::endl
                  << "r    : 刷新帮助信息" << std::endl
                  << "q/ESC: 退出" << std::endl
                  << "===================" << std::endl;
        refresh.store(false);
    }
}

// ------------- 主程序 -------------
int main()
{
    XInitThreads();

    try {
        // 获取输入PCD文件列表
        std::vector<fs::path> pcd_files;
        for (const auto& entry : fs::directory_iterator(INPUT_FOLDER)) {
            if (entry.path().extension() == ".pcd") pcd_files.push_back(entry.path());
        }
        if (pcd_files.empty()) throw std::runtime_error("No .pcd files found in directory: " + INPUT_FOLDER);
        std::sort(pcd_files.begin(), pcd_files.end());

        // 创建窗口
        pcl::visualization::PCLVisualizer::Ptr viewer(new pcl::visualization::PCLVisualizer("3D LiDAR Viewer"));
        viewer->setBackgroundColor(0.1, 0.1, 0.1);
        viewer->registerKeyboardCallback(keyboardEventOccurred, (void*)viewer.get());
        viewer->setCameraPosition(0, 0, 50, 0, 0, 1, 0, 1, 0);

        while (!viewer->wasStopped() && !quit.load()) {
            if (current_idx >= pcd_files.size()) current_idx = pcd_files.size() - 1;
            pcl::PointCloud<pcl::PointXYZI>::Ptr cloud(new pcl::PointCloud<pcl::PointXYZI>);
            pcl::PointCloud<pcl::PointXYZI>::Ptr ground(new pcl::PointCloud<pcl::PointXYZI>);
            pcl::PointCloud<pcl::PointXYZI>::Ptr obstacles(new pcl::PointCloud<pcl::PointXYZI>);

            // 1 加载
            if (pcl::io::loadPCDFile<pcl::PointXYZI>(pcd_files[current_idx].string(), *cloud) == -1) {
                std::cerr << "Error loading: " << pcd_files[current_idx] << std::endl;
                continue;
            }

            // 2 地面分离
            ground_remove::removeGround(cloud, ground, obstacles);

            // 3 聚类
            std::array<std::array<int, numGrid>, numGrid> mask;
            int numCluster = 0;
            componentClustering(obstacles, mask, numCluster);

            std::vector<pcl::PointCloud<pcl::PointXYZI>> objects;
            getClusteredPoints(obstacles, mask, objects);

            std::vector<pcl::PointCloud<pcl::PointXYZI>> bbPoints;
            boxFitting(objects, bbPoints);

            std::cout << "聚类数: " << numCluster << std::endl;
            for (size_t i = 0; i < objects.size(); ++i) {
                std::cout << "Cluster " << (i+1) << " 点数: " << objects[i].size() << std::endl;
            }

            // 4 可视化
            viewer->removeAllPointClouds();
            viewer->removeAllShapes();

            // 地面点云绿色
            pcl::visualization::PointCloudColorHandlerCustom<pcl::PointXYZI> ground_color(ground, GROUND_R, GROUND_G, GROUND_B);
            viewer->addPointCloud<pcl::PointXYZI>(ground, ground_color, "ground");
            viewer->setPointCloudRenderingProperties(pcl::visualization::PCL_VISUALIZER_POINT_SIZE, 2, "ground");

            // 聚类+检测框 可视化
            for (size_t i = 0; i < objects.size(); ++i) {
                if (objects[i].empty()) continue;
                auto color = generateClusterColor(i, objects.size());
                std::string obj_name = "cluster_" + std::to_string(i);
                pcl::visualization::PointCloudColorHandlerCustom<pcl::PointXYZI>
                    cluster_color(objects[i].makeShared(), std::get<0>(color), std::get<1>(color), std::get<2>(color));
                viewer->addPointCloud<pcl::PointXYZI>(objects[i].makeShared(), cluster_color, obj_name);
                viewer->setPointCloudRenderingProperties(pcl::visualization::PCL_VISUALIZER_POINT_SIZE, 5, obj_name);

                // 检测框
                if (i < bbPoints.size() && bbPoints[i].size() == 8) {
                    visualizeBoundingBox(viewer, bbPoints[i], color, "box_" + std::to_string(i));
                }
            }

            // // 5 保存地面和地上
            // saveGround(ground, current_idx);
            // saveObjects(objects, current_idx);

            // 帮助提示
            displayControlInfo();

            viewer->spinOnce(50);

            // 控制
            while (!viewer->wasStopped() && !quit.load() && !next_flag.load() && !prev_flag.load()) {
                viewer->spinOnce(50);
                std::this_thread::sleep_for(std::chrono::milliseconds(50));
            }
            if (next_flag.load()) {
                if (current_idx + 1 < pcd_files.size()) ++current_idx;
                next_flag.store(false);
            } else if (prev_flag.load()) {
                if (current_idx > 0) --current_idx;
                prev_flag.store(false);
            }
        }
    } catch (const std::exception& e) {
        std::cerr << "运行出错: " << e.what() << std::endl;
        return EXIT_FAILURE;
    }

    std::cout << "\nProgram terminated successfully." << std::endl;
    return EXIT_SUCCESS;
}